"""
视频流处理模块
支持RTMP和RTSP协议的视频流接收和帧处理功能
"""

import cv2
import numpy as np
from PyQt5.QtCore import QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QImage, QPixmap
import time
import threading


class VideoStream(QThread):
    """统一的视频流处理类，支持RTMP和RTSP协议"""
    
    # 信号定义
    frame_ready = pyqtSignal(np.ndarray, np.ndarray)  # BGR帧, HSV帧
    connection_status = pyqtSignal(str)  # 连接状态消息
    error_occurred = pyqtSignal(str)  # 错误信息
    
    def __init__(self):
        super().__init__()
        self.stream_url = None
        self.protocol = None  # 'rtmp' or 'rtsp'
        self.cap = None
        self.is_running = False
        self.fps = 30  # 默认帧率
        self.frame_skip = 1  # 帧跳跃，用于控制处理频率
        self.frame_count = 0
        
    def set_stream_url(self, url):
        """设置流地址并自动检测协议"""
        self.stream_url = url
        if url:
            url_lower = url.lower()
            if url_lower.startswith('rtmp://'):
                self.protocol = 'rtmp'
            elif url_lower.startswith('rtsp://'):
                self.protocol = 'rtsp'
            else:
                self.protocol = None
                
    def get_protocol(self):
        """获取当前协议类型"""
        return self.protocol
        
    def set_fps_limit(self, fps):
        """设置帧率限制"""
        self.fps = max(1, min(fps, 60))
        
    def set_frame_skip(self, skip):
        """设置帧跳跃数"""
        self.frame_skip = max(1, skip)
        
    def start_stream(self):
        """启动视频流"""
        if not self.stream_url:
            self.error_occurred.emit("请先设置视频流地址")
            return False
            
        if not self.protocol:
            self.error_occurred.emit("不支持的协议，仅支持RTMP和RTSP")
            return False
            
        try:
            protocol_name = self.protocol.upper()
            self.connection_status.emit(f"正在连接{protocol_name}流...")
            
            # 创建VideoCapture对象
            self.cap = cv2.VideoCapture(self.stream_url)
            
            # 对于RTSP流，设置额外的参数以减少延迟
            if self.protocol == 'rtsp':
                self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
                # 设置RTSP传输协议，优先使用TCP
                self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('H','2','6','4'))
            else:
                # RTMP流设置
                self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            
            # 检查连接是否成功
            if not self.cap.isOpened():
                self.error_occurred.emit(f"无法连接到{protocol_name}流，请检查地址是否正确")
                return False
                
            # 获取流信息
            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            stream_fps = self.cap.get(cv2.CAP_PROP_FPS)
            
            self.connection_status.emit(f"{protocol_name}流连接成功 - 分辨率: {width}x{height}, FPS: {stream_fps:.1f}")
            
            self.is_running = True
            self.start()  # 启动线程
            return True
            
        except Exception as e:
            protocol_name = self.protocol.upper() if self.protocol else "视频"
            self.error_occurred.emit(f"连接{protocol_name}流时发生错误: {str(e)}")
            return False
            
    def stop_stream(self):
        """停止视频流"""
        self.is_running = False
        if self.cap is not None:
            self.cap.release()
            self.cap = None
        protocol_name = self.protocol.upper() if self.protocol else "视频"
        self.connection_status.emit(f"{protocol_name}流已断开")
        
    def run(self):
        """线程主循环"""
        frame_interval = 1.0 / self.fps  # 帧间隔
        last_frame_time = time.time()
        
        while self.is_running and self.cap is not None:
            try:
                ret, frame = self.cap.read()
                
                if not ret:
                    protocol_name = self.protocol.upper() if self.protocol else "视频"
                    self.error_occurred.emit(f"读取{protocol_name}流帧失败，可能是网络中断")
                    break
                    
                current_time = time.time()
                
                # 帧率控制
                if current_time - last_frame_time < frame_interval:
                    continue
                    
                # 帧跳跃控制
                self.frame_count += 1
                if self.frame_count % self.frame_skip != 0:
                    continue
                    
                # 转换颜色空间
                frame_hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
                
                # 发射信号
                self.frame_ready.emit(frame, frame_hsv)
                
                last_frame_time = current_time
                
            except Exception as e:
                protocol_name = self.protocol.upper() if self.protocol else "视频"
                self.error_occurred.emit(f"处理{protocol_name}流时发生错误: {str(e)}")
                break
                
        self.stop_stream()


# 为了向后兼容，保留RTMPStream类
class RTMPStream(VideoStream):
    """RTMP流处理类（向后兼容）"""
    
    def __init__(self):
        super().__init__()
        self.rtmp_url = None
        
    def set_rtmp_url(self, url):
        """设置RTMP流地址（向后兼容方法）"""
        self.rtmp_url = url
        self.set_stream_url(url)


class VideoStreamManager:
    """视频流管理器"""
    
    def __init__(self):
        self.stream_thread = None
        self.current_frame = None
        self.current_frame_hsv = None
        self.frame_callbacks = []  # 帧回调函数列表
        
    def create_stream(self):
        """创建新的视频流对象"""
        if self.stream_thread is not None:
            self.stop_stream()
            
        self.stream_thread = VideoStream()
        self.stream_thread.frame_ready.connect(self._on_frame_received)
        return self.stream_thread
        
    def _on_frame_received(self, frame_bgr, frame_hsv):
        """处理接收到的帧"""
        self.current_frame = frame_bgr.copy()
        self.current_frame_hsv = frame_hsv.copy()
        
        # 调用所有回调函数
        for callback in self.frame_callbacks:
            try:
                callback(frame_bgr, frame_hsv)
            except Exception as e:
                print(f"帧回调函数执行错误: {str(e)}")
                
    def add_frame_callback(self, callback):
        """添加帧回调函数"""
        if callback not in self.frame_callbacks:
            self.frame_callbacks.append(callback)
            
    def remove_frame_callback(self, callback):
        """移除帧回调函数"""
        if callback in self.frame_callbacks:
            self.frame_callbacks.remove(callback)
            
    def get_current_frame(self):
        """获取当前帧"""
        return self.current_frame, self.current_frame_hsv
        
    def start_stream(self, stream_url, fps_limit=30, frame_skip=1):
        """启动视频流"""
        if self.stream_thread is None:
            self.create_stream()
            
        self.stream_thread.set_stream_url(stream_url)
        self.stream_thread.set_fps_limit(fps_limit)
        self.stream_thread.set_frame_skip(frame_skip)
        
        return self.stream_thread.start_stream()
        
    def stop_stream(self):
        """停止视频流"""
        if self.stream_thread is not None:
            self.stream_thread.stop_stream()
            self.stream_thread.wait()  # 等待线程结束
            self.stream_thread = None
            
        self.current_frame = None
        self.current_frame_hsv = None
        
    def is_streaming(self):
        """检查是否正在流式传输"""
        return self.stream_thread is not None and self.stream_thread.is_running
        
    def get_protocol(self):
        """获取当前协议类型"""
        if self.stream_thread is not None:
            return self.stream_thread.get_protocol()
        return None
        
    def get_stream_info(self):
        """获取流信息"""
        if self.stream_thread is None or self.current_frame is None:
            return None
            
        height, width = self.current_frame.shape[:2]
        return {
            'width': width,
            'height': height,
            'stream_url': self.stream_thread.stream_url,
            'protocol': self.stream_thread.protocol,
            'fps_limit': self.stream_thread.fps,
            'frame_skip': self.stream_thread.frame_skip
        }


# 为了向后兼容，保留RTMPStreamManager类
class RTMPStreamManager(VideoStreamManager):
    """RTMP流管理器（向后兼容）"""
    
    def create_stream(self):
        """创建新的RTMP流对象（向后兼容）"""
        if self.stream_thread is not None:
            self.stop_stream()
            
        self.stream_thread = RTMPStream()
        self.stream_thread.frame_ready.connect(self._on_frame_received)
        return self.stream_thread
        
    def start_stream(self, rtmp_url, fps_limit=30, frame_skip=1):
        """启动RTMP流（向后兼容方法）"""
        return super().start_stream(rtmp_url, fps_limit, frame_skip)
        
    def get_stream_info(self):
        """获取流信息（向后兼容）"""
        info = super().get_stream_info()
        if info:
            # 添加rtmp_url字段以保持向后兼容
            info['rtmp_url'] = info['stream_url']
        return info


def validate_stream_url(url):
    """验证视频流地址格式"""
    if not url:
        return False, "视频流地址不能为空"
        
    url_lower = url.lower()
    
    if url_lower.startswith('rtmp://'):
        return validate_rtmp_url(url)
    elif url_lower.startswith('rtsp://'):
        return validate_rtsp_url(url)
    else:
        return False, "不支持的协议，仅支持RTMP和RTSP协议"


def validate_rtmp_url(url):
    """验证RTMP地址格式"""
    if not url:
        return False, "RTMP地址不能为空"
        
    if not url.lower().startswith('rtmp://'):
        return False, "RTMP地址必须以 'rtmp://' 开头"
        
    # 基本格式检查
    try:
        parts = url.replace('rtmp://', '').split('/')
        if len(parts) < 2:
            return False, "RTMP地址格式不正确，应为: rtmp://server/app/stream"
        return True, "RTMP地址格式正确"
    except:
        return False, "RTMP地址格式验证失败"


def validate_rtsp_url(url):
    """验证RTSP地址格式"""
    if not url:
        return False, "RTSP地址不能为空"
        
    if not url.lower().startswith('rtsp://'):
        return False, "RTSP地址必须以 'rtsp://' 开头"
        
    # 基本格式检查
    try:
        # RTSP URL格式通常为: rtsp://[username:password@]host[:port]/path
        url_part = url.replace('rtsp://', '')
        
        # 检查是否包含必要的主机部分
        if '/' in url_part:
            host_part = url_part.split('/')[0]
        else:
            host_part = url_part
            
        if not host_part:
            return False, "RTSP地址格式不正确，缺少主机地址"
            
        return True, "RTSP地址格式正确"
    except:
        return False, "RTSP地址格式验证失败"


def convert_frame_to_qpixmap(frame_bgr):
    """将OpenCV帧转换为Qt Pixmap"""
    if frame_bgr is None:
        return None
        
    try:
        # BGR转RGB
        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        height, width, channel = frame_rgb.shape
        bytes_per_line = 3 * width
        
        # 创建QImage
        q_image = QImage(frame_rgb.data, width, height, bytes_per_line, QImage.Format_RGB888)
        
        # 转换为QPixmap
        pixmap = QPixmap.fromImage(q_image)
        return pixmap
        
    except Exception as e:
        print(f"帧转换错误: {str(e)}")
        return None